#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import numpy as np
from scene import Scene
import os
from tqdm import tqdm
from os import makedirs
from gaussian_renderer import render
import torchvision
from utils.general_utils import safe_state, getProjectionMatrix
from argparse import ArgumentParser
from arguments import ModelParams, PipelineParams, get_combined_args
from gaussian_renderer import GaussianModel
import copy, time
import imageio

gen_image = lambda x: (x[0] * x[1] + x[2]) * x[3] 
gen_sun_shade = lambda x: x[0] * x[1]
gen_shade = lambda x: x[0] * x[1] + x[2]
to8b = lambda x: x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()


def render_set(model_path, name, iteration, views, gaussians, pipeline, background):
    render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "renders")
    gts_path = os.path.join(model_path, name, "ours_{}".format(iteration), "gt")
    d_render_path = os.path.join(model_path, name, "ours_{}".format(iteration), "de_renders")

    makedirs(render_path, exist_ok=True)
    makedirs(gts_path, exist_ok=True)
    makedirs(d_render_path, exist_ok=True)

    t_list = []
    for idx, view in enumerate(tqdm(views, desc="Rendering progress")):
        if view.split == name:
            
            torch.cuda.synchronize(); t0 =time.time()
            renderings = render(view, gaussians, pipeline, background)["render"]
            torch.cuda.synchronize();t1= time.time()
            t_list.append(t1-t0)
            
            rendering = gen_image(renderings)
            mask = view.mask.cuda()
            gt = view.original_image[0:3, :, :].cuda()
            de_rendering = torch.cat([torch.cat(renderings, -1), 
                                      torch.cat([mask[:,:,None].repeat(1,1,3), renderings[0]*renderings[1], rendering, gt], -1)], 1)
            torchvision.utils.save_image(de_rendering, os.path.join(d_render_path, '{0:05d}'.format(view.uid) + ".png"))
            
            if view.split == 'eval':
                H, W = rendering.shape[1:]
                rendering, gt = rendering[:, :, W//2:], gt[:, :, W//2:]
            torchvision.utils.save_image(rendering, os.path.join(render_path, '{0:05d}'.format(view.uid) + ".png"))
            torchvision.utils.save_image(gt, os.path.join(gts_path, '{0:05d}'.format(view.uid) + ".png"))
            
    t = np.array(t_list[5:])
    fps = 1.0 / t.mean()
    print(f'Test FPS: \033[1;35m{fps:.5f}\033[0m')

def interpolate_views(view0, view1, length):
    #view0 source view
    #view1 target view 
    results = []
    for i in range(length):
        weight = (i+1) / length
        view = copy.deepcopy(view0)
        # view.camera_center = view0.camera_center * (1-weight) + view1.camera_center * weight
        view.FoVy = view0.FoVy * (1-weight) + view1.FoVy * weight
        view.FoVx = view0.FoVx * (1-weight) + view1.FoVx * (i/length)
        view.projection_matrix = view0.projection_matrix * (1-weight) + view1.projection_matrix * weight
        view.world_view_transform = view0.world_view_transform * (1-weight) + view1.world_view_transform * weight
        view.full_proj_transform = (view.world_view_transform.unsqueeze(0).bmm(view.projection_matrix.unsqueeze(0))).squeeze(0)
        results.append(view)
    return results

def generate_multi_views(views, length=60):
    generated_views = []
    for i in range(len(views)-1):
        views_temp = interpolate_views(views[i], views[i+1], length)
        generated_views.extend(views_temp)
    return generated_views

def render_vedio(model_path, train_views, gaussians, pipeline, background):
    scene_name = model_path.split('/')[-1]
    if scene_name == "brandenburg_gate":
        format_idx = 11
        select_view_id = [12, 59, 305]
        length_view = 90 * 2
        vis_idxs = [53, 33, 41, 73, 84, 101]
        sun_idxs = [53, 33, 41, 73, 84, 101]
        env_idxs = [53, 33, 41, 73, 84, 101]
        views = [train_views[i] for i in select_view_id]
        
    elif scene_name == "sacre_coeur":
        format_idx = 38
        select_view_id = [753, 657, 595, 181, 699] # 700
        length_view = 45 * 2
        vis_idxs = [350, 76]
        sun_idxs = [350, 76]
        env_idxs = [350, 76]
        views=[train_views[i] for i in select_view_id]
        
    elif scene_name == "trevi_fountain":
        format_idx = 17
        select_view_id = [408, 303, 79, 893, 395, 281] # 700 需要插帧的基础视角
        length_view = 45 * 2    # 视频总帧数
        vis_idxs = [317,495]  # 采用的外观
        sun_idxs = [317,495]  # 采用的外观
        env_idxs = [317,495]  # 采用的外观
        views = [train_views[i] for i in select_view_id]
        
    generated_views = generate_multi_views(views, length_view)
    
    video_images_path = os.path.join(model_path, f"video_images") #
    makedirs(video_images_path, exist_ok=True)
    
    render_video_out = imageio.get_writer(f'{video_images_path}/{scene_name}_{vis_idxs[0]}' + '.mp4', mode='I', fps=30,codec='libx264',quality=10.0)#
    
    n_views = len(generated_views)
    for idx, view in enumerate(tqdm(generated_views, desc="Rendering progress")):
        view.image_height, view.image_width = train_views[format_idx].image_height, train_views[format_idx].image_width
        
        # interpolate visibility
        # vis_total = n_views // (len(vis_idxs) - 1)
        # weight = idx % vis_total / vis_total
        # vis_idx0 = vis_idxs[idx // vis_total]
        # vis_idx1 = vis_idxs[idx // vis_total + 1]
        # vis_config = {'idx0': vis_idx0, 'idx1': vis_idx1, 'weight': weight}
        
        # sun_total = n_views // (len(sun_idxs) - 1)
        # weight = idx % sun_total / sun_total
        # sun_idx0 = sun_idxs[idx // sun_total]
        # sun_idx1 = sun_idxs[idx // sun_total + 1]
        # sun_config = {'idx0': sun_idx0, 'idx1': sun_idx1, 'weight': weight}
        
        # env_total = n_views // (len(env_idxs) - 1)
        # weight = idx % env_total / env_total
        # env_idx0 = env_idxs[idx // env_total]
        # env_idx1 = env_idxs[idx // env_total + 1]
        # env_config = {'idx0': env_idx0, 'idx1': env_idx1, 'weight': weight}
        
        vis_config = vis_idxs[0]
        sun_config = sun_idxs[0]
        env_config = env_idxs[0]
        
        
        renderings = render(view, gaussians, pipeline, background, vis_config, sun_config, env_config)["render"]
        rendered = gen_image(renderings)
        sun_shade = gen_sun_shade(renderings)
        shade = gen_shade(renderings)
        tem_image = torch.cat([torch.ones_like(shade), sun_shade, shade, rendered], -1)
        
        image = torch.cat([torch.cat(renderings, -1), tem_image], 1)
        np_image = to8b(image)
        render_video_out.append_data(np_image)
        torchvision.utils.save_image(image, os.path.join(video_images_path, f"video_image_{idx:05d}.png"))
        
        del renderings, image, np_image  
    render_video_out.close()

def render_sets(dataset: ModelParams, iteration: int, pipeline: PipelineParams, skip_train: bool, skip_test: bool, render_interpolate: bool, render_multiview_vedio: bool):
    with torch.no_grad():
        gaussians = GaussianModel(dataset)
        scene = Scene(dataset, gaussians, load_iteration=iteration, shuffle=False)

        bg_color = [1,1,1] if dataset.white_background else [0, 0, 0]
        background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
        model_path = dataset.model_path

        # if not skip_train:
        #      render_set(model_path, "train", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background)

        # if not skip_test:
        #      render_set(model_path, "eval", scene.loaded_iter, scene.getTrainCameras(), gaussians, pipeline, background)
        
        if not render_multiview_vedio:
            render_vedio(model_path, scene.getTrainCameras(), gaussians, pipeline, background)
        

if __name__ == "__main__":
    # Set up command line argument parser
    parser = ArgumentParser(description="Testing script parameters")
    model = ModelParams(parser, sentinel=True)
    pipeline = PipelineParams(parser)
    parser.add_argument("--iteration", default=-1, type=int)
    parser.add_argument("--skip_train", action="store_true")
    parser.add_argument("--skip_test", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    
    parser.add_argument("--render_interpolate", action="store_true", default=False)
    parser.add_argument("--render_multiview_vedio", action="store_true", default=False)
    
    args = get_combined_args(parser)
    print("Rendering " + args.model_path)

    # Initialize system state (RNG)
    safe_state(args.quiet)

    render_sets(model.extract(args), 
                args.iteration, 
                pipeline.extract(args), 
                args.skip_train, 
                args.skip_test,
                args.render_interpolate,
                args.render_multiview_vedio)